import { randomUUID } from 'node:crypto';
import type { WritableStream } from 'node:stream/web';
import type { RequestContext } from '../../di';
import { ErrorDomain, ErrorCategory } from '../../error';
import type { MastraScorers } from '../../evals';
import { runScorer } from '../../evals/hooks';
import { SpanType, wrapMastra } from '../../observability';
import type { TracingContext } from '../../observability';
import type { ChunkType } from '../../stream/types';
import { ToolStream } from '../../tools/stream';
import type { DynamicArgument } from '../../types';
import { EMITTER_SYMBOL, STREAM_FORMAT_SYMBOL } from '../constants';
import type { DefaultExecutionEngine } from '../default';
import type { Step, SuspendOptions } from '../step';
import { getStepResult } from '../step';
import type {
  Emitter,
  ExecutionContext,
  RestartExecutionParams,
  SerializedStepFlowEntry,
  StepExecutionResult,
  StepResult,
  TimeTravelExecutionParams,
} from '../types';
import {
  validateStepInput,
  createDeprecationProxy,
  runCountDeprecationMessage,
  validateStepResumeData,
  validateStepSuspendData,
  validateStepStateData,
} from '../utils';

export interface ExecuteStepParams {
  workflowId: string;
  runId: string;
  resourceId?: string;
  step: Step<string, any, any>;
  stepResults: Record<string, StepResult<any, any, any, any>>;
  executionContext: ExecutionContext;
  restart?: RestartExecutionParams;
  timeTravel?: TimeTravelExecutionParams;
  resume?: {
    steps: string[];
    resumePayload: any;
    label?: string;
    forEachIndex?: number;
  };
  prevOutput: any;
  emitter: Emitter;
  abortController: AbortController;
  requestContext: RequestContext;
  skipEmits?: boolean;
  writableStream?: WritableStream<ChunkType>;
  disableScorers?: boolean;
  serializedStepGraph: SerializedStepFlowEntry[];
  tracingContext: TracingContext;
  iterationCount?: number;
}

export async function executeStep(
  engine: DefaultExecutionEngine,
  params: ExecuteStepParams,
): Promise<StepExecutionResult> {
  const {
    workflowId,
    runId,
    resourceId,
    step,
    stepResults,
    executionContext,
    restart,
    resume,
    timeTravel,
    prevOutput,
    emitter,
    abortController,
    requestContext,
    skipEmits = false,
    writableStream,
    disableScorers,
    serializedStepGraph,
    tracingContext,
    iterationCount,
  } = params;

  const stepCallId = randomUUID();

  const { inputData, validationError } = await validateStepInput({
    prevOutput,
    step,
    validateInputs: engine.options?.validateInputs ?? true,
  });

  const { resumeData: timeTravelResumeData, validationError: timeTravelResumeValidationError } =
    await validateStepResumeData({
      resumeData: timeTravel?.stepResults[step.id]?.status === 'suspended' ? timeTravel?.resumeData : undefined,
      step,
    });

  let resumeDataToUse: unknown;
  if (timeTravelResumeData && !timeTravelResumeValidationError) {
    resumeDataToUse = timeTravelResumeData;
  } else if (timeTravelResumeData && timeTravelResumeValidationError) {
    engine.getLogger().warn('Time travel resume data validation failed', {
      stepId: step.id,
      error: timeTravelResumeValidationError.message,
    });
  } else if (resume?.steps[0] === step.id) {
    resumeDataToUse = resume?.resumePayload;
  }

  // Extract suspend data if this step was previously suspended
  let suspendDataToUse =
    stepResults[step.id]?.status === 'suspended' ? stepResults[step.id]?.suspendPayload : undefined;

  // Filter out internal workflow metadata before exposing to step code
  if (suspendDataToUse && '__workflow_meta' in suspendDataToUse) {
    const { __workflow_meta, ...userSuspendData } = suspendDataToUse;
    suspendDataToUse = userSuspendData;
  }

  const startTime = resumeDataToUse ? undefined : Date.now();
  const resumeTime = resumeDataToUse ? Date.now() : undefined;

  const stepInfo = {
    ...stepResults[step.id],
    ...(resumeDataToUse ? { resumePayload: resumeDataToUse } : { payload: inputData }),
    ...(startTime ? { startedAt: startTime } : {}),
    ...(resumeTime ? { resumedAt: resumeTime } : {}),
    status: 'running',
    ...(iterationCount ? { metadata: { iterationCount } } : {}),
  };

  executionContext.activeStepsPath[step.id] = executionContext.executionPath;

  const stepSpan = tracingContext.currentSpan?.createChildSpan({
    name: `workflow step: '${step.id}'`,
    type: SpanType.WORKFLOW_STEP,
    input: inputData,
    attributes: {
      stepId: step.id,
    },
    tracingPolicy: engine.options?.tracingPolicy,
  });

  const operationId = `workflow.${workflowId}.run.${runId}.step.${step.id}.running_ev`;
  await engine.onStepExecutionStart({
    step,
    inputData,
    emitter,
    executionContext,
    stepCallId,
    stepInfo,
    operationId,
    skipEmits,
  });

  await engine.persistStepUpdate({
    workflowId,
    runId,
    resourceId,
    serializedStepGraph,
    stepResults: {
      ...stepResults,
      [step.id]: stepInfo,
    } as Record<string, StepResult<any, any, any, any>>,
    executionContext,
    workflowStatus: 'running',
    requestContext,
  });

  // Check if this is a nested workflow that requires special handling
  if (engine.isNestedWorkflowStep(step)) {
    const workflowResult = await engine.executeWorkflowStep({
      step,
      stepResults,
      executionContext,
      resume,
      timeTravel,
      prevOutput,
      inputData,
      emitter,
      startedAt: startTime ?? Date.now(),
      abortController,
      requestContext,
      tracingContext,
      writableStream,
      stepSpan,
    });

    // If executeWorkflowStep returns a result, wrap it in StepExecutionResult
    if (workflowResult !== null) {
      const stepResult = { ...stepInfo, ...workflowResult } as StepResult<any, any, any, any>;
      return {
        result: stepResult,
        stepResults: { [step.id]: stepResult },
        mutableContext: engine.buildMutableContext(executionContext),
        requestContext: engine.serializeRequestContext(requestContext),
      };
    }
  }

  const runStep = async (data: any) => {
    // Wrap data with a Proxy to show deprecation warning for runCount
    const proxiedData = createDeprecationProxy(data, {
      paramName: 'runCount',
      deprecationMessage: runCountDeprecationMessage,
      logger: engine.getLogger(),
    });

    return step.execute(proxiedData);
  };

  let execResults: any;

  const retries = step.retries ?? executionContext.retryConfig.attempts ?? 0;
  const delay = executionContext.retryConfig.delay ?? 0;

  // Use executeStepWithRetry to handle retry logic
  // Default engine: internal retry loop
  // Inngest engine: throws RetryAfterError for external retry handling
  const stepRetryResult = await engine.executeStepWithRetry(
    `workflow.${workflowId}.step.${step.id}`,
    async () => {
      if (validationError) {
        throw validationError;
      }

      const retryCount = engine.getOrGenerateRetryCount(step.id);

      let timeTravelSteps: string[] = [];
      if (timeTravel && timeTravel.steps.length > 0) {
        timeTravelSteps = timeTravel.steps[0] === step.id ? timeTravel.steps.slice(1) : [];
      }

      let suspended: { payload: any } | undefined;
      let bailed: { payload: any } | undefined;
      const contextMutations: {
        suspendedPaths: Record<string, number[]>;
        resumeLabels: Record<string, { stepId: string; foreachIndex?: number }>;
        stateUpdate: any;
        requestContextUpdate: Record<string, any> | null;
      } = {
        suspendedPaths: {},
        resumeLabels: {},
        stateUpdate: null,
        requestContextUpdate: null,
      };

      const output = await runStep({
        runId,
        resourceId,
        workflowId,
        mastra: engine.mastra ? wrapMastra(engine.mastra, { currentSpan: stepSpan }) : undefined,
        requestContext,
        inputData,
        state: executionContext.state,
        setState: async (state: any) => {
          const { stateData, validationError: stateValidationError } = await validateStepStateData({
            stateData: state,
            step,
            validateInputs: engine.options?.validateInputs ?? true,
          });
          if (stateValidationError) {
            throw stateValidationError;
          }
          // executionContext.state = stateData;
          contextMutations.stateUpdate = stateData;
        },
        retryCount,
        resumeData: resumeDataToUse,
        suspendData: suspendDataToUse,
        tracingContext: { currentSpan: stepSpan },
        getInitData: () => stepResults?.input as any,
        getStepResult: getStepResult.bind(null, stepResults),
        suspend: async (suspendPayload?: any, suspendOptions?: SuspendOptions): Promise<any> => {
          const { suspendData, validationError: suspendValidationError } = await validateStepSuspendData({
            suspendData: suspendPayload,
            step,
            validateInputs: engine.options?.validateInputs ?? true,
          });
          if (suspendValidationError) {
            throw suspendValidationError;
          }
          // Capture mutations for return value (needed for Inngest replay)
          contextMutations.suspendedPaths[step.id] = executionContext.executionPath;
          // Also apply directly for Default engine
          executionContext.suspendedPaths[step.id] = executionContext.executionPath;

          if (suspendOptions?.resumeLabel) {
            const resumeLabel = Array.isArray(suspendOptions.resumeLabel)
              ? suspendOptions.resumeLabel
              : [suspendOptions.resumeLabel];
            for (const label of resumeLabel) {
              const labelData = {
                stepId: step.id,
                foreachIndex: executionContext.foreachIndex,
              };
              // Capture for return value
              contextMutations.resumeLabels[label] = labelData;
              // Apply directly for Default engine
              executionContext.resumeLabels[label] = labelData;
            }
          }

          suspended = { payload: suspendData };
        },
        bail: (result: any) => {
          bailed = { payload: result };
        },
        abort: () => {
          abortController?.abort();
        },
        // Only pass resume data if this step was actually suspended before
        // This prevents pending nested workflows from trying to resume instead of start
        resume:
          stepResults[step.id]?.status === 'suspended'
            ? {
                steps: resume?.steps?.slice(1) || [],
                resumePayload: resume?.resumePayload,
                // @ts-ignore
                runId: stepResults[step.id]?.suspendPayload?.__workflow_meta?.runId,
                label: resume?.label,
                forEachIndex: resume?.forEachIndex,
              }
            : undefined,
        // Only pass restart data if this step is part of activeStepsPath
        // This prevents pending nested workflows from trying to restart instead of start
        restart: !!restart?.activeStepsPath?.[step.id],
        timeTravel:
          timeTravelSteps.length > 0
            ? {
                inputData: timeTravel?.inputData,
                steps: timeTravelSteps,
                nestedStepResults: timeTravel?.nestedStepResults,
                resumeData: timeTravel?.resumeData,
              }
            : undefined,
        [EMITTER_SYMBOL]: emitter,
        [STREAM_FORMAT_SYMBOL]: executionContext.format,
        engine: engine.getEngineContext(),
        abortSignal: abortController?.signal,
        writer: new ToolStream(
          {
            prefix: 'workflow-step',
            callId: stepCallId,
            name: step.id,
            runId,
          },
          writableStream,
        ),
        // Disable scorers must be explicitly set to false they are on by default
        scorers: disableScorers === false ? undefined : step.scorers,
        validateInputs: engine.options?.validateInputs,
      });

      // Capture requestContext state after step execution (only for engines that need it)
      if (engine.requiresDurableContextSerialization()) {
        contextMutations.requestContextUpdate = engine.serializeRequestContext(requestContext);
      }

      return { output, suspended, bailed, contextMutations };
    },
    { retries, delay, stepSpan, workflowId, runId },
  );

  // Check if step execution failed
  if (!stepRetryResult.ok) {
    execResults = stepRetryResult.error;
  } else {
    const { result: durableResult } = stepRetryResult;

    // Apply context mutations from the durable operation result
    // For Default: these were already applied during execution, this is a no-op
    // For Inngest: on replay, the wrapped function didn't re-execute, so we restore from the memoized result
    Object.assign(executionContext.suspendedPaths, durableResult.contextMutations.suspendedPaths);
    Object.assign(executionContext.resumeLabels, durableResult.contextMutations.resumeLabels);

    // Restore requestContext from memoized result (only for engines that need it)
    if (engine.requiresDurableContextSerialization() && durableResult.contextMutations.requestContextUpdate) {
      requestContext.clear();
      for (const [key, value] of Object.entries(durableResult.contextMutations.requestContextUpdate)) {
        requestContext.set(key, value);
      }
    }

    if (step.scorers) {
      await runScorersForStep({
        engine,
        scorers: step.scorers,
        runId,
        input: inputData,
        output: durableResult.output,
        workflowId,
        stepId: step.id,
        requestContext,
        disableScorers,
        tracingContext: { currentSpan: stepSpan },
      });
    }

    if (durableResult.suspended) {
      execResults = {
        status: 'suspended',
        suspendPayload: durableResult.suspended.payload,
        ...(durableResult.output ? { suspendOutput: durableResult.output } : {}),
        suspendedAt: Date.now(),
      };
    } else if (durableResult.bailed) {
      execResults = { status: 'bailed', output: durableResult.bailed.payload, endedAt: Date.now() };
    } else {
      execResults = { status: 'success', output: durableResult.output, endedAt: Date.now() };
    }
  }

  delete executionContext.activeStepsPath[step.id];

  if (!skipEmits) {
    const emitOperationId = `workflow.${workflowId}.run.${runId}.step.${step.id}.emit_result`;
    await engine.wrapDurableOperation(emitOperationId, async () => {
      await emitStepResultEvents({
        stepId: step.id,
        stepCallId,
        execResults: { ...stepInfo, ...execResults } as StepResult<any, any, any, any>,
        emitter,
      });
    });
  }

  if (execResults.status != 'failed') {
    stepSpan?.end({
      output: execResults.output,
      attributes: {
        status: execResults.status,
      },
    });
  }

  const stepResult = { ...stepInfo, ...execResults } as StepResult<any, any, any, any>;

  return {
    result: stepResult,
    stepResults: { [step.id]: stepResult },
    mutableContext: engine.buildMutableContext({
      ...executionContext,
      state: stepRetryResult.ok
        ? (stepRetryResult.result.contextMutations.stateUpdate ?? executionContext.state)
        : executionContext.state,
    }),
    requestContext: engine.serializeRequestContext(requestContext),
  };
}

export interface RunScorersParams {
  engine: DefaultExecutionEngine;
  scorers: DynamicArgument<MastraScorers>;
  runId: string;
  input: any;
  output: any;
  requestContext: RequestContext;
  workflowId: string;
  stepId: string;
  disableScorers?: boolean;
  tracingContext: TracingContext;
}

export async function runScorersForStep(params: RunScorersParams): Promise<void> {
  const { engine, scorers, runId, input, output, workflowId, stepId, requestContext, disableScorers, tracingContext } =
    params;

  let scorersToUse = scorers;
  if (typeof scorersToUse === 'function') {
    try {
      scorersToUse = await scorersToUse({
        requestContext: requestContext,
      });
    } catch (error) {
      engine.preprocessExecutionError(
        error,
        {
          id: 'WORKFLOW_FAILED_TO_FETCH_SCORERS',
          domain: ErrorDomain.MASTRA_WORKFLOW,
          category: ErrorCategory.USER,
          details: {
            runId,
            workflowId,
            stepId,
          },
        },
        'Error fetching scorers: ',
      );
    }
  }

  if (!disableScorers && scorersToUse && Object.keys(scorersToUse || {}).length > 0) {
    for (const [_id, scorerObject] of Object.entries(scorersToUse || {})) {
      runScorer({
        scorerId: scorerObject.name,
        scorerObject: scorerObject,
        runId: runId,
        input: input,
        output: output,
        requestContext,
        entity: {
          id: workflowId,
          stepId: stepId,
        },
        structuredOutput: true,
        source: 'LIVE',
        entityType: 'WORKFLOW',
        tracingContext,
      });
    }
  }
}

/**
 * Emit step result events (suspended, result, finish).
 * Shared between Default and Inngest execution engines.
 */
export async function emitStepResultEvents(params: {
  stepId: string;
  stepCallId?: string;
  execResults: StepResult<any, any, any, any>;
  emitter: Emitter;
}): Promise<void> {
  const { stepId, stepCallId, execResults, emitter } = params;
  const payloadBase = stepCallId ? { id: stepId, stepCallId } : { id: stepId };

  if (execResults.status === 'suspended') {
    await emitter.emit('watch', {
      type: 'workflow-step-suspended',
      payload: { ...payloadBase, ...execResults },
    });
  } else {
    await emitter.emit('watch', {
      type: 'workflow-step-result',
      payload: { ...payloadBase, ...execResults },
    });
    await emitter.emit('watch', {
      type: 'workflow-step-finish',
      payload: { ...payloadBase, metadata: {} },
    });
  }
}
